%given referenced data, returns training and testing sets using different
%strategies.
% function [training testing] = getTrainingTesting(history, ntraining, ntesting, ntransition)
% params: 
% - history: data collected in the form of cells of trajectories (each
%            trajectory itself is a cell of transitions. If the data we
%            have is in the form of histories, we should first convert them
%            into one flat history before calling this function. 
% - ntraining: number of training sample points to return. 
% - ntesting:  number of testing sample points to return. 
% - ntransition: number of transitions to be collected from each trajectory
% -------------------------------------------------------------------------

function [training testing] = getTrainingTesting(history, ntraining, ntesting, ntransition)

if nargin < 4 
   ntransition = -1; %collect all transitions  
end

%set testing data on a uniform grid.
%  d = floor(sqrt(ntesting)); 
%  testing = discretize(STATE_SPEC, [d,d]); 
%  testing = [testing (floor(rand(d*d,1)*4)+1)]; %add random action

testing = collectTesting(history, ntesting); 

training = collectTraining(history, ntraining, ntransition); 
end

function testing = collectTesting(history, ntesting)
inds = floor ( rand(ntesting,1)*length(history)+1) ; 

inds = sort(inds); 

testing = zeros(ntesting, 2); 
for i=1:ntesting 
    testing(i,:) = history{inds(i)}{1}.state; 
end
end



%returns the first ntransition element of each trajectory inside history
%until ntraining of them is collected. if ntransition ==-1 , all of the
%transitions are collected from each trajectory. 
function training = collectTraining(history, ntraining, ntransition)

cnt = ntraining; 

st = history{1}{1}; 
if (ntraining==-1)
   ntraining = 0; 
   for i=1:length(history)
       ntraining = ntraining + length(history{i});
   end
end

training = repmat(st, 1, ntraining); 

hpos = 1; 
tpos = 1; 
for i=1:ntraining
    training(i) = history{hpos}{tpos}; 
    if  (tpos >=ntransition && ntransition ~= -1) ||  tpos >= length(history{hpos}) %should move to the next trajectory
        tpos = 1; 
        hpos = hpos +1; 
        if hpos > length(history)
            if (cnt==-1)
                training(i+1:length(training))= []; 
            else
                display('ERROR: not enough data to generate training testing'); 
            end
            return; 
        end
    else %still in the middle of current trajectory
       tpos = tpos + 1;  
    end
end
end

